package utils

import (
	"bytes"
	"crypto/des"
	"errors"
)

func PKCS5Padding(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padText...)
}

func PKCS5UnPadding(src []byte) ([]byte, error) {
	if len(src) == 0 {
		return nil, errors.New("padding is illegal")
	}
	padding := int(src[len(src)-1])
	if padding < 1 || padding > des.BlockSize {
		return nil, errors.New("padding is illegal")
	}
	for i := len(src) - padding; i < len(src); i++ {
		if int(src[i]) != padding {
			return nil, errors.New("padding is illegal")
		}
	}
	return src[:len(src)-padding], nil
}

func DesEncode(key []byte, plaintext []byte) ([]byte, error) {
	block, err := des.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(plaintext)%des.BlockSize != 0 {
		plaintext = PKCS5Padding(plaintext, des.BlockSize)
	}

	ciphertext := make([]byte, len(plaintext))
	for i := 0; i < len(plaintext); i += des.BlockSize {
		block.Encrypt(ciphertext[i:i+des.BlockSize], plaintext[i:i+des.BlockSize])
	}

	return ciphertext, nil
}

func DesDecode(key []byte, ciphertext []byte) ([]byte, error) {
	block, err := des.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(ciphertext)%des.BlockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	plaintext := make([]byte, len(ciphertext))
	for i := 0; i < len(ciphertext); i += des.BlockSize {
		block.Decrypt(plaintext[i:i+des.BlockSize], ciphertext[i:i+des.BlockSize])
	}

	unpadText, err := PKCS5UnPadding(plaintext)
	if err != nil {
		return nil, err
	}

	return unpadText, nil
}
